# -*- coding: utf-8 -*-
"""
3D 曲面图组件
"""
from .echarts_base import TgEchartsBase
from pyecharts.charts import Surface3D
from pyecharts import options as opts
import pandas as pd
import numpy as np

class TgSurface3DChart(TgEchartsBase):
    """3D 曲面图组件"""

    def __init__(self, parent=None):
        super().__init__(parent)
        self.x_column = None
        self.y_column = None
        self.z_column = None
        self.chart_title = "3D 曲面图"
        self.color_map = None

    def set_data(self, df: pd.DataFrame, x_column: str, y_column: str, z_column: str):
        """
        设置3D曲面图数据
        :param df: 数据DataFrame
        :param x_column: X轴对应的数据列名
        :param y_column: Y轴对应的数据列名
        :param z_column: Z轴对应的数据列名
        """
        self.df = df
        self.x_column = x_column
        self.y_column = y_column
        self.z_column = z_column
        self.update_chart()

    def set_title(self, title: str):
        """设置图表标题"""
        self.chart_title = title
        if hasattr(self, 'df') and self.df is not None:
            self.update_chart()

    def set_color_map(self, color_map: str = 'viridis'):
        """
        设置颜色映射
        :param color_map: 颜色映射名称，可选: 'viridis', 'plasma', 'cool', 'hot'
        """
        self.color_map = color_map
        if hasattr(self, 'df') and self.df is not None:
            self.update_chart()

    def _prepare_data(self):
        """准备3D曲面数据"""
        if (self.df is None or 
            self.df.empty or 
            self.x_column not in self.df.columns or 
            self.y_column not in self.df.columns or 
            self.z_column not in self.df.columns):
            return [], [], []
        
        # 获取唯一的x和y值并排序
        x_unique = sorted(self.df[self.x_column].unique())
        y_unique = sorted(self.df[self.y_column].unique())
        
        # 创建网格数据
        z_data = []
        for y_val in y_unique:
            row = []
            for x_val in x_unique:
                # 查找对应的z值
                mask = (self.df[self.x_column] == x_val) & (self.df[self.y_column] == y_val)
                if mask.any():
                    z_val = self.df.loc[mask, self.z_column].iloc[0]
                else:
                    z_val = np.nan
                row.append(z_val)
            z_data.append(row)
        
        # 转换为Surface3D需要的格式
        surface_data = []
        for i, y_val in enumerate(y_unique):
            for j, x_val in enumerate(x_unique):
                if not np.isnan(z_data[i][j]):
                    surface_data.append([j, i, z_data[i][j]])
        
        return x_unique, y_unique, surface_data

    def create_chart(self) -> Surface3D:
        """创建3D曲面图"""
        if self.df is None or self.x_column is None or self.y_column is None or self.z_column is None:
            chart = Surface3D()
            chart.set_global_opts(title_opts=opts.TitleOpts(title=self.chart_title))
            return chart

        x_unique, y_unique, surface_data = self._prepare_data()
        
        chart = Surface3D()
        
        # 添加曲面数据
        chart.add(
            series_name="Surface",
            data=surface_data,
            xaxis3d_opts=opts.Axis3DOpts(
                name=self.x_column,
                type_="category",
                data=[str(x) for x in x_unique]
            ),
            yaxis3d_opts=opts.Axis3DOpts(
                name=self.y_column,
                type_="category", 
                data=[str(y) for y in y_unique]
            ),
            zaxis3d_opts=opts.Axis3DOpts(
                name=self.z_column,
                type_="value"
            ),
            grid3d_opts=opts.Grid3DOpts(
                width=100,
                height=100,
                depth=100,
                is_rotate=False,  # 禁用自动旋转
                rotate_speed=0,   # 自动旋转速度为0
                rotate_sensitivity=1,  # 保持鼠标拖拽敏感度
            ),
        )
        
        # 设置全局选项
        chart.set_global_opts(
            title_opts=opts.TitleOpts(title=self.chart_title),
            visualmap_opts=opts.VisualMapOpts(
                min_=min([item[2] for item in surface_data]) if surface_data else 0,
                max_=max([item[2] for item in surface_data]) if surface_data else 100,
                range_color=self._get_color_range(),
                is_calculable=True,
                dimension=2,
                pos_left="0",
                pos_top="50%",
            ),
            toolbox_opts=opts.ToolboxOpts(
                feature={
                    "saveAsImage": {},
                    "restore": {},
                    "dataZoom": {},
                }
            ),
        )
        
        return chart

    def _get_color_range(self):
        """获取颜色范围"""
        color_maps = {
            'viridis': ["#440154", "#31688e", "#35b779", "#fde725"],
            'plasma': ["#0d0887", "#7e03a8", "#cc4778", "#f89441", "#f0f921"],
            'cool': ["#00ffff", "#0080ff", "#0000ff", "#8000ff", "#ff00ff"],
            'hot': ["#000000", "#800000", "#ff0000", "#ff8000", "#ffff00", "#ffffff"],
            'default': ["#313695", "#4575b4", "#74add1", "#abd9e9", "#e0f3f8", 
                       "#ffffcc", "#fee090", "#fdae61", "#f46d43", "#d73027", "#a50026"]
        }
        return color_maps.get(self.color_map, color_maps['default'])

    def update_chart(self):
        """更新图表"""
        if self.df is None:
            return
        
        chart = self.create_chart()
        self.render_chart(chart)

    def clear(self):
        """Clear chart content and reset internal data/state."""
        if hasattr(self, 'webview') and self.webview is not None:
            try:
                self.webview.setHtml("<html><head><meta charset='utf-8'></head><body></body></html>")
            except Exception:
                pass
        self.chart = None
        self.df = None
        self.x_column = None
        self.y_column = None
        self.z_column = None

    def set_grid3d_options(self, width=100, height=100, depth=100, is_rotate=False, rotate_speed=0, rotate_sensitivity=1):
        """
        设置3D网格选项
        :param width: 网格宽度
        :param height: 网格高度  
        :param depth: 网格深度
        :param is_rotate: 是否启用自动旋转
        :param rotate_speed: 自动旋转速度
        :param rotate_sensitivity: 鼠标拖拽旋转敏感度
        """
        self.grid3d_opts = {
            'width': width,
            'height': height,
            'depth': depth,
            'is_rotate': is_rotate,
            'rotate_speed': rotate_speed,
            'rotate_sensitivity': rotate_sensitivity
        }
        if hasattr(self, 'df') and self.df is not None:
            self.update_chart()
